#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import asyncio
import json
import sys
from typing import List

from src.agents.AbstractGeneratorAgent import AbstractGeneratorAgent
from src.agents.KeywordExtractorAgent import KeywordExtractorAgent
from src.agents.LiteratureRetrieverAgent import LiteratureRetrieverAgent
from src.agents.OutlineBuilderAgent import OutlineBuilderAgent
from src.agents.base_config import StateManager
from src.utils.env_config import get_task_name


class TaskOrchestrator:
    """主控智能体：负责任务编排和流程控制"""

    def __init__(self):
        self.state_manager = StateManager()
        self.keyword_agent = KeywordExtractorAgent()
        self.literature_agent = LiteratureRetrieverAgent()
        self.abstract_agent = AbstractGeneratorAgent()
        self.outline_agent = OutlineBuilderAgent()

    async def execute_task(self, topic: str) -> None:
        """执行完整的报告生成任务"""
        print(f"开始执行报告生成任务：{topic}")

        # 保存主题到上下文
        self.state_manager.save_context("topic", topic)

        # T1: 关键词提取
        await self._execute_subtask("T1", self._task1_keyword_extraction, topic)

        # T2: 文献检索
        keywords = self.state_manager.get_context("keywords")
        await self._execute_subtask("T2", self._task2_literature_retrieval, keywords)

        # T3: 摘要生成
        retrieval_summary = self.state_manager.get_context("retrieval_summary")
        await self._execute_subtask("T3", self._task3_abstract_generation, retrieval_summary)

        # T4: 报告框架搭建
        abstract = self.state_manager.get_context("abstract")
        await self._execute_subtask("T4", self._task4_outline_building, topic, abstract)

        print("所有任务执行完成！")
        print(f"状态文件：{self.state_manager.state_file}")
        print("报告框架文件：report_framework.md")

    async def _execute_subtask(self, task_id: str, task_func, *args) -> None:
        """执行子任务"""
        task_name = get_task_name(task_id)
        print(f"\n开始执行 {task_id}: {task_name}")

        # 更新状态为"进行中"
        self.state_manager.update_task_status(task_id, "进行中")
        try:
            # 执行任务
            await task_func(*args)
            # 更新状态为"完成"
            self.state_manager.update_task_status(task_id, "完成")
            print(f"{task_id}: {task_name} 执行完成")

        except Exception as e:
            print(f"{task_id}: {task_name} 执行失败: {e}")
            self.state_manager.update_task_status(task_id, "失败")
            raise

    async def _task1_keyword_extraction(self, topic: str) -> None:
        """T1: 关键词提取"""
        keywords = await self.keyword_agent.extract_keywords(topic)
        self.state_manager.save_context("keywords", keywords)
        print(f"提取的关键词: {keywords}")

    async def _task2_literature_retrieval(self, keywords: List[str]) -> None:
        """T2: 文献检索"""
        retrieval_summary = await self.literature_agent.retrieve_literature(keywords)
        self.state_manager.save_context("retrieval_summary", retrieval_summary)
        print(f"文献检索摘要: {retrieval_summary[:100]}...")

    async def _task3_abstract_generation(self, retrieval_summary: str) -> None:
        """T3: 摘要生成"""
        abstract = await self.abstract_agent.generate_abstract(retrieval_summary)
        self.state_manager.save_context("abstract", abstract)
        print(f"生成的摘要: {abstract[:100]}...")

    async def _task4_outline_building(self, topic: str, abstract: str) -> None:
        """T4: 报告框架搭建"""
        markdown_content = await self.outline_agent.build_outline(topic, abstract)

        # 保存到文件
        with open("report_framework.md", "w", encoding="utf-8") as f:
            f.write(markdown_content)

        self.state_manager.save_context("markdown_content", markdown_content)
        print("报告框架已生成并保存到 report_framework.md")


async def main():
    # """主函数"""
    # if len(sys.argv) != 2:
    #     print("使用方法: python task1_main.py '报告主题'")
    #     print("示例: python task1_agentscope.py '2025 Q3 AI 行业趋势'")
    #     return
    # topic = sys.argv[1]
    topic = '2025 Q3 AI 行业趋势'
    # 创建任务编排器并执行任务
    orchestrator = TaskOrchestrator()
    await orchestrator.execute_task(topic)


if __name__ == "__main__":
    asyncio.run(main())